# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags
from concurrent.futures import ProcessPoolExecutor

import collections
import datetime
import functools
import logging
import os

import numpy
import pandas
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

import scipy.stats

from coin.base.datetime_util import to_timestamp_int

from coin.base.book.types import TradeSide
import coin.strategy.mm.tool.archive_base as abase


def fill_zero(seri):
  plt.fill_between(
      pandas.Series(seri).index,
      seri * 0,
      seri,
      where=seri < 0,
      facecolor='red',
      alpha=0.5,
      lw=0.2,
      rasterized=True)
  plt.fill_between(
      pandas.Series(seri).index,
      seri * 0,
      seri,
      where=seri > 0,
      facecolor='green',
      alpha=0.5,
      lw=0.2,
      rasterized=True)


class FeedArchiveChecker2(object):
  def __init__(self, products, start_ts=None, end_ts=None, min_hilo_diff=None):
    self._products = products
    self._tick_sizes = (0.5, 0.01)
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)
    self._fields_out = []
    self._fields_view = []
    self._time_set = set()
    self._agg_book = {}
    self._tsigned_cum_qty = collections.defaultdict(float)
    self._trades = {product: collections.OrderedDict() for product in products}
    self._ftbooks = {product: collections.OrderedDict() for product in products}
    self._etbooks = {product: collections.OrderedDict() for product in products}
    self._cmbooks = {product: collections.OrderedDict() for product in products}
    self._et_cm_map = {}
    self._cmetbooks = {product: collections.OrderedDict() for product in products}
    self._instrs = {product: collections.OrderedDict() for product in products}
    self._min_hilo_diff = 0 if min_hilo_diff is None else float(min_hilo_diff)
    self._prev_midp = None
    self._ft_in_unit = None
    self._et_in_unit = None
    self._cum_movement = 0

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      if hasattr(book_builder_name, "exchange"):
        if book_builder_name.exchange != product.exchange:
          continue

      if not hasattr(book_builder, "ProductType") or isinstance(product, book_builder.ProductType):
        book_builder.subscribe(product, functools.partial(self.on_book, product))
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

        # print(f"+sub {product}:{book_builder_name}")
        if hasattr(book_builder, "subscribe_instrument"):
          book_builder.subscribe_instrument(product, functools.partial(self.on_instrument, product))

  def on_instrument(self, product, instrument):
    if instrument.timestamp < self._start_ts:
      return

    if self._end_ts and self._end_ts < instrument.timestamp:
      return

    if hasattr(instrument, 'open_interest'):
      open_interest = instrument.open_interest
      funding_rate = instrument.funding_rate
      mark_price = instrument.mark_price
      if hasattr(instrument, 'indicative_funding_rate'):
        indicative_funding_rate = instrument.indicative_funding_rate
      else:
        indicative_funding_rate = None
    elif product.exchange == "Bitmex":
      open_interest = instrument.inst.get('openInterest', None)
      funding_rate = instrument.inst.get('fundingRate', None)
      mark_price = instrument.inst.get('markPrice', None)
      indicative_funding_rate = instrument.inst.get('indicativeFundingRate', None)
    else:
      print(instrument)
      return

    ts_in_unit = int(instrument.timestamp * 1e-9) * 1e9

    self._instrs[product][ts_in_unit] = {
        'ts_in_unit': ts_in_unit,
        'timestamp_instr': instrument.timestamp,
        'open_interest': open_interest,
        'funding_rate': funding_rate,
        'indicative_funding_rate': indicative_funding_rate,
        'mark_price': mark_price,
    }

  def on_trade(self, product, trade):
    if trade.timestamp < self._start_ts:
      return

    if self._end_ts and self._end_ts < trade.timestamp:
      return

    ts_in_unit = int(trade.timestamp * 1e-9) * 1e9

    tradesign = 1 if trade.side == TradeSide.TRADE_BUY_SIDE else -1

    self._tsigned_cum_qty[product] += trade.qty * tradesign

    self._trades[product][ts_in_unit] = {
        'ts_in_unit': ts_in_unit, 'tradecumsq': self._tsigned_cum_qty[product]
    }

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return

    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
      if product == self._products[0]:
        midp = 0.5 * (bid0_price + ask0_price)
        if self._prev_midp is not None:
          retbp = (midp / self._prev_midp - 1) * 1e4
          self._cum_movement += abs(retbp)
        self._prev_midp = midp
    except (AttributeError, IndexError):
      return

    unit = 1e6
    if product == self._products[0]:
      self._ft_in_unit = int(book.timestamp / unit) * unit
      self._et_in_unit = int(book.exchange_timestamp / unit) * unit

    if self._ft_in_unit is None:
      return

    self._ftbooks[product][self._ft_in_unit] = {
        'ts_in_unit': self._ft_in_unit,
        'timestamp': book.timestamp,
        'bid0p': bid0_price,
        'ask0p': ask0_price
    }

    noself_et_in_unit = int(book.exchange_timestamp / unit) * unit

    self._etbooks[product][self._et_in_unit] = {
        'ts_in_unit': self._et_in_unit,
        'timestamp': book.exchange_timestamp,
        'noself_et_in_unit': noself_et_in_unit,
        'bid0p': bid0_price,
        'ask0p': ask0_price
    }

    cm_in_unit = int(self._cum_movement / 1) * 1
    self._cmbooks[product][cm_in_unit] = {
        'ts_in_unit': cm_in_unit,
        'timestamp': book.timestamp,
        'bid0p': bid0_price,
        'ask0p': ask0_price
    }
    self._et_cm_map[self._et_in_unit] = cm_in_unit


  def plot_all(self, out_dir, postfix, trading_date_str, machine):
    for product, cmetbook in self._cmetbooks.items():
      if len(cmetbook) == 0:
        for _, item in self._etbooks[product].items():
          et_in_unit = item['noself_et_in_unit']
          if et_in_unit in self._et_cm_map:
            cm_in_unit = self._et_cm_map[et_in_unit]
            cmetbook[cm_in_unit] = item.copy()
            cmetbook[cm_in_unit]['ts_in_unit'] = cm_in_unit

    from matplotlib.backends.backend_pdf import PdfPages
    with PdfPages(f"{out_dir}/{trading_date_str}_{postfix}_{machine}.pdf") as pdf:
      i = 0
      product1 = list(self._ftbooks.keys())[0]
      for j, product2 in enumerate(self._ftbooks.keys()):
        if i >= j:
          continue
        self.plot_product_comp(
            trading_date_str,
            product1,
            product2,
            "[FT]" + machine + postfix,
            self._ftbooks,
            "fetched time",
            pandas.DatetimeIndex,
            out_dir,
            pdf)
        self.plot_product_comp(
            trading_date_str,
            product1,
            product2,
            "[ET]" + machine + postfix,
            self._etbooks,
            "exchange time",
            pandas.DatetimeIndex,
            out_dir,
            pdf)
        self.plot_product_comp(
            trading_date_str,
            product1,
            product2,
            "[MOVEBPS]" + machine + postfix,
            self._cmbooks,
            "cum midp move bps",
            numpy.array,
            out_dir,
            pdf)
        self.plot_product_comp(
            trading_date_str,
            product1,
            product2,
            "[MOVEBPS_ET]" + machine + postfix,
            self._cmetbooks,
            "[BOTH_ET]cum midp move bps",
            numpy.array,
            out_dir,
            pdf)

  def plot_product_comp(self, date_str, product1, product2, namestr, books, xlabel, ttype, out_dir, pdf):
    if product1 not in self._products or product2 not in self._products:
      return None, None, None
    if len(books[product1]) == 0 or len(books[product2]) == 0:
      return None, None, None
    legends = []
    dfs = []

    for product, col, suffix in [(product1, 'g', 'x'), (product2, 'r', 'y')]:
      df = pandas.DataFrame(list(books[product].values()))
      df = df.sort_values('ts_in_unit').fillna(method='ffill').reset_index(drop=True)
      timeindex = ttype(df['ts_in_unit'])

      df = df.rename(
          columns={col: f'{col}_{suffix}' for col in df.columns if col not in 'ts_in_unit'})
      dfs.append(df)
      legends.append(f'{product.exchange}_{product.relative_norm}')

    ddf = pandas.merge(
        dfs[0], on='ts_in_unit', right=dfs[1],
        how='outer').sort_values('ts_in_unit').fillna(method='ffill').reset_index(drop=True)

    ddf['mid0p_x'] = 0.5 * (ddf['ask0p_x'] + ddf['bid0p_x'])
    ddf['mid0p_y'] = 0.5 * (ddf['ask0p_y'] + ddf['bid0p_y'])
    ddf['mid0p_diff_xy'] = ddf['mid0p_x'] - ddf['mid0p_y']

    plt.plot(
        ttype(ddf['ts_in_unit']),
        ddf['ask0p_x'],
        'g-',
        ttype(ddf['ts_in_unit']),
        ddf['ask0p_y'],
        'r-',
        ttype(ddf['ts_in_unit']),
        numpy.nan * ddf['bid0p_x'],
        'b-',
        ttype(ddf['ts_in_unit']),
        ddf['bid0p_x'],
        'g-',
        ttype(ddf['ts_in_unit']),
        ddf['bid0p_y'],
        'r-',
        lw=0.2,
        markersize=0.2,
        drawstyle='steps-post',
        rasterized=True)

    def axvlines(xs, *args, **kwargs):
      for x in xs:
        plt.axvline(x, *args, **kwargs)


    if ttype == numpy.array:
      tss = pandas.DatetimeIndex(ddf['timestamp_x'])
      idxseri = pandas.Series(ddf['ts_in_unit'])
      idxseri.index = tss
      tot_time = tss.max() - tss.min()
      xtick_unit = tot_time / 10
      idxdf = idxseri.resample(xtick_unit).first().reset_index()
      idxdf = idxdf.loc[~numpy.isnan(idxdf['ts_in_unit'])]
      plt.xticks(
          idxdf['ts_in_unit'].astype(int),
          pandas.DatetimeIndex(idxdf['timestamp_x']).strftime("%m%d %H:%M:%S"),
          rotation=90)
      axvlines(idxdf['ts_in_unit'].astype(int), color='y', lw=0.2)
    plt.tight_layout()
    plt.xlabel(xlabel)
    legends.append(f'{product2.exchange}_{product2.relative_norm}-{product1.exchange}_{product1.relative_norm}')
    plt.legend(legends, loc='upper right')
    plt.twinx()

    ddf['mid0p_diff_xy'] = -ddf['mid0p_diff_xy']

    plt.plot(
        ttype(ddf['ts_in_unit']),
        ddf['mid0p_diff_xy'],
        'b-',
        ttype(ddf['ts_in_unit']),
        ddf['mid0p_diff_xy'].rolling(300).mean(),
        'k-',
        lw=0.2,
        alpha=0.5,
        drawstyle='steps-post',
        rasterized=True)

    signal = ddf['mid0p_diff_xy'] - ddf['mid0p_diff_xy'].rolling(300).mean()
    plt.twinx()
    signal.index = ttype(ddf['ts_in_unit'])
    fill_zero(signal)

    title = "%s_%s_%s-%s" % (date_str, namestr, legends[0], legends[1])
    plt.title(f'{title}')

    #plt.minorticks_on()
    #plt.grid(lw=0.5, which='major')
    #plt.grid(lw=0.2, which='minor')
    plt.tight_layout()
    filepath = os.path.join(out_dir, "%s.png" % title)
    print(filepath)
    pdf.savefig()
    plt.savefig(filepath, dpi=300)
    plt.close()
    return (legends[0], legends[1])


def launch(baseflags, fsflags, symbols, subsymbols, min_hilo_diff, out_dir):
  plt.rcParams['figure.figsize'] = 8, 6
  plt.rcParams['font.size'] = 6
  plt.rcParams['legend.fontsize'] = 6
  plt.rcParams['xtick.labelsize'] = 6
  plt.rcParams['ytick.labelsize'] = 6

  products = []
  if symbols is not None:
    products.extend(abase.symbols_to_products(symbols, baseflags))

  if subsymbols is not None:
    products.extend(abase.subsymbols_to_products(subsymbols, baseflags))

  start_time, end_time = abase.get_time_range(baseflags)

  strategy = FeedArchiveChecker2(products=products,
                                 start_ts=start_time,
                                 end_ts=end_time,
                                 min_hilo_diff=min_hilo_diff)

  abase.run_from_archive_base(baseflags=baseflags,
                              feed_sub_flags_proto=fsflags,
                              products=products,
                              on_book_reset=strategy.on_book_reset)

  strategy.plot_all(out_dir,
                    postfix=flags.FLAGS.postfix,
                    trading_date_str=baseflags.trading_date,
                    machine=fsflags.feed_machine)


def launch_date(baseflags, fsflags, out_dir, symbols, subsymbols, min_hilo_diff):
  launch(baseflags, fsflags, symbols, subsymbols, min_hilo_diff, out_dir)


def main(argv):
  if flags.FLAGS.symbols is None and flags.FLAGS.subsymbols is None:
    flags.FLAGS.subsymbols = (
        'Okex:Futures:BTC-USD.QUARTER,Huobi:Futures:BTC-USD.QUARTER'
        + ',Okex:Futures:BTC-USD.PERPETUAL,Binance:Futures:BTC-USDT.PERPETUAL'
        + ',Bitmex:Futures:BTC-USD.PERPETUAL')
  if flags.FLAGS.api_override is None:
    flags.FLAGS.api_override = (
        'Futures:Okex=v3.realtime_light_bbo'
        + ',Futures:Huobi=v1.realtime_light_bbo'
        + ',Futures:Binance=v1.l1_realtime'
        + ',Futures:Ftx=v1.realtime'
        + ',Futures:Bitmex=v1.realtime_light_bbo')

  symbols = None if flags.FLAGS.symbols is None else flags.FLAGS.symbols.split(",")
  subsymbols = None if flags.FLAGS.subsymbols is None else flags.FLAGS.subsymbols.split(",")

  jobs = []

  flags.FLAGS.out_dir = os.path.join(flags.FLAGS.out_dir, datetime.datetime.now().strftime("%Y%m%d_%H%M%S"))
  os.makedirs(flags.FLAGS.out_dir, exist_ok=True)
  print("http://files.ap.ngrok.io" + os.path.realpath(flags.FLAGS.out_dir))

  for trading_date in abase.get_trading_dates(flags.FLAGS.trading_date):
    flags.FLAGS.trading_date = trading_date.strftime("%Y%m%d")
    baseflags = abase.get_base_flags()
    fsflags = abase.get_feed_subscription_flags()
    jobs.append((launch_date,
                 baseflags,
                 fsflags,
                 flags.FLAGS.out_dir,
                 symbols,
                 subsymbols,
                 flags.FLAGS.min_hilo_diff))

  if len(jobs) == 1:
    jobs[0][0](*jobs[0][1:])
  else:
    futures = []
    with ProcessPoolExecutor(max_workers=8) as executor:
      for job in jobs:
        futures.append(executor.submit(job[0], *job[1:]))
      for future in futures:
        future.result()


# oper
# plot_leadlag --postfix=major_btc --subsymbols=Binance:Futures:BTC-USDT.PERPETUAL,Okex:Futures:BTC-USD.QUARTER,Huobi:Futures:BTC-USD.QUARTER,Okex:Futures:BTC-USD.PERPETUAL --trading_date=20211204-20211204 --time_range=4-6

if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws.huobi'
  flags.DEFINE_string(
      'symbols',
      None,
      'symbol name')
  flags.DEFINE_string('out_dir', 'pic_ll', 'Output directory.')
  flags.DEFINE_string('postfix', None, 'filename postfix')
  flags.DEFINE_string(
      'subsymbols',
      None,
      'subscript symbol name')
  flags.DEFINE_string('min_hilo_diff', None, 'min hi-lo diff pct to pick')

  # --subsymbols: ex=Huobi:Futures:BTC-USD.QUARTER

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  app.run(main)
